﻿/* Copyright 2007-2008 dnAnalytics Project.
 *
 * Contributors to this file:
 * Jurgen Van Gael
 *
 * This file is part of dnAnalytics.  dnAnalytics is licensed under the 
 * Microsoft Public License. See License.txt for a complete copy of the
 * license.
 */
using System;
using dnAnalytics.LinearAlgebra;
using dnAnalytics.LinearAlgebra.Decomposition;
using dnAnalytics.Math;
using dnAnalytics.Properties;

namespace dnAnalytics.Statistics.Distributions
{
    /// <summary>
    /// This class implements functionality for the Wishart distribution. This distribution is
    /// parameterized by the degrees of freedom nu and the scale matrix S. The Wishart distribution
    /// is the conjugate prior for the precision (inverse covariance) matrix of the multivariate
    /// normal distribution.
    /// </summary>
    /// <remarks>The distribution will use the <see cref="System.Random"/> by default. 
    /// Users can set the random number generator by using the <see cref="RandomNumberGenerator"/> property.
    /// 
    /// The statistics classes will check all the incoming parameters whether they are in the allowed
    /// range. This might involve heavy computation. Optionally, by setting Control.CheckDistributionParameters
    /// to false, all parameter checks can be turned off.</remarks>
    public class Wishart
    {
        // The degrees of freedom for the Wishart distribution.
        private readonly double mNu;

        // The scale matrix for the Wishart distribution.
        private readonly Matrix mS;

        // Caches the Cholesky factorization of the scale matrix.
        private readonly Cholesky mChol;

        /// <summary>
        /// Construct a new Wishart distribution.
        /// </summary>
        /// <param name="nu">The degrees of freedom for the Wishart distribution.</param>
        /// <param name="S">The scale matrix for the Wishart distribution.</param>
        public Wishart(double nu, Matrix S)
        {
            mNu = nu;
            mS = S;
            mChol = new Cholesky(mS);

            if (Control.CheckDistributionParameters)
            {
                CheckParameters(nu, mChol);
            }

            RandomNumberGenerator = new System.Random();
        }

        /// <summary>
        /// The degrees of freedom for the Wishart distribution.
        /// </summary>
        public double Nu
        {
            get { return mNu; }
        }

        /// <summary>
        /// The scale matrix for the Wishart distribution.
        /// </summary>
        public Matrix S
        {
            get { return mS; }
        }

        /// <summary>
        /// Gets the mean.
        /// </summary>
        /// <value>The mean.</value>
        public Matrix Mean
        {
            get { return mNu*mS; }
        }

        /// <summary>
        /// Gets the mode.
        /// </summary>
        /// <value>The mode.</value>
        public Matrix Mode
        {
            get { return (mNu - mS.Rows - 1.0)*mS; }
        }

        /// <summary>
        /// Gets or sets the random number generator.
        /// </summary>
        /// <value>The random number generator used to generate a random sample.</value>
        public System.Random RandomNumberGenerator { get; set; }

        /// <summary>
        /// Checks the parameters of a Wishart distribution.
        /// </summary>
        /// <param name="nu">The degrees of freedom for the Wishart distribution.</param>
        /// <param name="S">The scale matrix for the Wishart distribution.</param>
        /// <exception cref="ArgumentOutOfRangeException">If the degrees of freedom parameter is negative.</exception>
        /// <exception cref="NotPositiveDefiniteException">If the scale parameter is not positive definite.</exception>
        public static void CheckParameters(double nu, Matrix S)
        {
            CheckParameters(nu, new Cholesky(S));
        }

        /// <summary>
        /// Checks the parameters of a Wishart distribution.
        /// </summary>
        /// <param name="nu">The degrees of freedom for the Wishart distribution.</param>
        /// <param name="chol">The cholesky factorization of the scale matrix for the Wishart distribution.</param>
        /// <exception cref="ArgumentOutOfRangeException">If the degrees of freedom parameter is negative.</exception>
        /// <exception cref="NotPositiveDefiniteException">If the scale parameter is not positive definite.</exception>
        public static void CheckParameters(double nu, Cholesky chol)
        {
            if (nu < 0.0)
            {
                throw new ArgumentOutOfRangeException("nu", Resources.NotPositive);
            }

            // Try a matrix factorization of S.
            if (!chol.IsPositiveDefinite())
            {
                throw new NotPositiveDefiniteException();
            }
        }

        /// <summary>
        /// Evaluates the probability density function for the Wishart distribution.
        /// </summary>
        /// <param name="X">The matrix at which to evaluate the density at.</param>
        /// <exception cref="ArgumentOutOfRangeException">If the argument does not have the same dimensions as the scale matrix.</exception>
        public double Density(Matrix X)
        {
            int p = mS.Rows;

            if (X.Rows != p || X.Columns != p)
            {
                throw new ArgumentOutOfRangeException("X", "X must be of same size as the scale matrix.");
            }

            double dX = X.Determinant();
            Matrix SiX = mChol.Solve(X);

            // Compute the multivariate Gamma function.
            double gp = System.Math.Pow(Math.Constants.Pi, p*(p - 1.0)/4.0);
            for (int j = 1; j <= p; j++)
            {
                gp *= SpecialFunctions.Gamma((mNu + 1.0 - j)/2.0);
            }

            return System.Math.Pow(dX, (mNu - p - 1.0)/2.0)
                   *System.Math.Exp(-0.5*SiX.Trace())
                   /System.Math.Pow(2.0, mNu*p/2.0)
                   /System.Math.Pow(mChol.Determinant(), mNu/2.0)
                   /gp;
        }

        /// <summary>
        /// Samples a Wishart distributed random variable using the method
        ///     Algorithm AS 53: Wishart Variate Generator
        ///     W. B. Smith and R. R. Hocking
        ///     Applied Statistics, Vol. 21, No. 3 (1972), pp. 341-345
        /// </summary>
        public Matrix Sample()
        {
            return DoSample(RandomNumberGenerator, mNu, mS, mChol);
        }

        /// <summary>
        /// Samples a Wishart distributed random variable using the method
        ///     Algorithm AS 53: Wishart Variate Generator
        ///     W. B. Smith and R. R. Hocking
        ///     Applied Statistics, Vol. 21, No. 3 (1972), pp. 341-345
        /// </summary>
        /// <param name="rnd">The random number generator to use.</param>
        /// <param name="nu">The degrees of freedom.</param>
        /// <param name="S">The scale matrix.</param>
        public static Matrix Sample(System.Random rnd, double nu, Matrix S)
        {
            if (Control.CheckDistributionParameters)
            {
                CheckParameters(nu, S);
            }

            return DoSample(rnd, nu, S, new Cholesky(S));
        }

        private static Matrix DoSample(System.Random rnd, double nu, Matrix S, Cholesky chol)
        {
            int D = S.Rows;

            // First generate a lower triangular matrix with Sqrt(Chi-Squares) on the diagonal
            // and normal distributed variables in the lower triangle.
            Matrix A = new DenseMatrix(D, D, 0.0);
            for (int d = 0; d < D; d++)
            {
                A[d, d] = System.Math.Sqrt(Gamma.Sample(rnd, (nu - d)/2.0, 0.5));
            }
            for (int i = 1; i < D; i++)
            {
                for (int j = 0; j < i; j++)
                {
                    A[i, j] = Normal.Sample(rnd, 0.0, 1.0);
                }
            }

            Matrix L = chol.Factor();

            // Now compute the Wishart distributed random matrix.
            return L*A*A.Transpose()*L.Transpose();
        }
    }
}